import os
import cv2
import numpy as np
from sklearn import svm
from sklearn.metrics import precision_score, recall_score, f1_score, roc_curve, auc
import time
from matplotlib import pyplot as plt
from sklearn.utils import shuffle
from imblearn.over_sampling import SMOTE
from tensorflow.keras.applications import VGG19
from tensorflow.keras.applications.vgg19 import preprocess_input
from tensorflow.keras.models import Model
from sklearn.model_selection import GridSearchCV
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 定义定位胸部的函数
def locate_chest(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    _, thresh = cv2.threshold(blurred, 45, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    max_contour = max(contours, key=cv2.contourArea)
    x, y, w, h = cv2.boundingRect(max_contour)
    cropped_image = image[y:y+h, x:x+w]
    return cropped_image

# 加载图像并应用胸部定位函数
def load_images_from_folder(folder):
    images = []
    labels = []
    for filename in os.listdir(folder):
        for img in os.listdir(os.path.join(folder, filename)):
            img_path = os.path.join(folder, filename, img)
            image = cv2.imread(img_path)
            if image is not None:
                image = locate_chest(image)
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                image = cv2.resize(image, (224, 224))
                images.append(image)
                labels.append(1 if filename == 'cancer' else 0)
    return np.array(images), np.array(labels)

# 数据增强
def augment_data(X, y):
    datagen = ImageDataGenerator(
        rotation_range=20,
        width_shift_range=0.2,
        height_shift_range=0.2,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True,
        fill_mode='nearest'
    )
    datagen.fit(X)
    return datagen

# 提取特征
def extract_features(model, images):
    images = preprocess_input(images)
    features = model.predict(images)
    return features

def train_model(X_train, y_train):
    param_grid = {
        'C': [0.1, 1, 10, 100],
        'gamma': [1, 0.1, 0.01, 0.001],
        'kernel': ['rbf', 'poly', 'sigmoid']
    }
    grid = GridSearchCV(svm.SVC(probability=True), param_grid, refit=True, verbose=2, cv=3)
    grid.fit(X_train, y_train)
    print(f"Best parameters found: {grid.best_params_}")
    return grid.best_estimator_

def predict(clf, X_test):
    y_pred = clf.predict(X_test)
    return y_pred

def print_metrics(y_test, y_pred):
    precision = precision_score(y_test, y_pred.round())
    recall = recall_score(y_test, y_pred.round())
    f1 = f1_score(y_test, y_pred.round())
    fpr, tpr, _ = roc_curve(y_test, y_pred)
    roc_auc = auc(fpr, tpr)
    print(f'Precision: {precision}\nRecall: {recall}\nF1 Score: {f1}\nAUC: {roc_auc}')
    plt.figure()
    plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % roc_auc)
    plt.plot([0, 1], [0, 1], 'k--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic')
    plt.legend(loc="lower right")
    plt.show()

def main():
    start_time = time.time()
    X_train, y_train = load_images_from_folder('D:\\kaggle\\final\\data1\\train')
    X_valid, y_valid = load_images_from_folder('D:\\kaggle\\final\\data1\\valid')
    X_test, y_test = load_images_from_folder('D:\\kaggle\\final\\data1\\test')
    X_train = np.concatenate((X_train, X_valid), axis=0)
    y_train = np.concatenate((y_train, y_valid), axis=0)
    X_train, y_train = shuffle(X_train, y_train, random_state=42)
    print(f'Training set size: {len(X_train)}, Test set size: {len(X_test)}')

    # 使用预训练的 VGG19 模型提取特征
    base_model = VGG19(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
    model = Model(inputs=base_model.input, outputs=base_model.output)
    X_train_features = extract_features(model, X_train)
    X_test_features = extract_features(model, X_test)
    X_train_features = X_train_features.reshape(X_train_features.shape[0], -1)
    X_test_features = X_test_features.reshape(X_test_features.shape[0], -1)

    # Use SMOTE to balance the data
    sm = SMOTE(random_state=42)
    X_train_features, y_train = sm.fit_resample(X_train_features, y_train)

    # 数据增强
    datagen = augment_data(X_train, y_train)

    clf = train_model(X_train_features, y_train)
    y_pred = predict(clf, X_test_features)
    print_metrics(y_test, y_pred)
    end_time = time.time()
    print(f'Program running time: {end_time - start_time}')

if __name__ == "__main__":
    main()